package io.gitee.luoyuan800.tool.llm.agent;

import io.gitee.luoyuan800.tool.llm.bean.AbstractMessage;
import io.gitee.luoyuan800.tool.llm.bean.ContentMessage;
import io.gitee.luoyuan800.tool.llm.bean.completion.YuanQiChatCompletion;
import io.gitee.luoyuan800.tool.llm.bean.config.AiAgentConfig;
import io.gitee.luoyuan800.tool.llm.bean.request.ChatRequest;
import io.gitee.luoyuan800.tool.llm.bean.request.YuanQiRequest;
import io.gitee.luoyuan800.tool.llm.web.RestfulClient;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;

import java.time.ZonedDateTime;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * 可以调用元器的智能体。
 * 当前版本只支持文本消息。
 */
@Slf4j
public class YuanQiAgent extends CommonAgent {
    private static final String URL = "https://open.hunyuan.tencent.com/openapi/v1/agent/chat/completions";

    public YuanQiAgent(AiAgentConfig aiAgentConfig, RestfulClient restTemplate, String name) {
        super(aiAgentConfig, restTemplate, name);
        if (StringUtils.isBlank(aiAgentConfig.getUrl())) {
            aiAgentConfig.getModelConfig().setFullUrl(URL);
        }
    }

    @Override
    YuanQiRequest buildRequest(int resultCount) {
        YuanQiRequest request = new YuanQiRequest();
        setRequestParams(resultCount, request);
        return request;
    }

    @Override
    void setRequestParams(int resultCount, ChatRequest request) {
        super.setRequestParams(resultCount, request);
        YuanQiRequest yuanQiRequest = (YuanQiRequest) request;
        yuanQiRequest.setAssistantId(getAiAgentConfig().getAssistantId());
        yuanQiRequest.setUserId(getAiAgentConfig().getUserId());
    }

    @Override
    AbstractMessage<?> buildMessage(ChatRequest.MessageType messageType, String msg, String image) {
        return new ContentMessage(msg, image);
    }

    @Override
    List<AbstractMessage<?>> addMessage(ChatRequest.MessageType messageType, String msg, String image) {
        while (getMessages().size() >= 39) {
            //最多支持40条历史消息
            getMessages().remove(0);
        }
        return super.addMessage(messageType, msg, image);
    }

    @Override
    public YuanQiChatCompletion fullChat(String msg, int resultCount, String image, boolean enableThinking) {
        log.info("Sending message to {}.", getName());
        log.debug("Message: {}", msg);

        YuanQiRequest request = buildRequest(resultCount);
        request.setStream(false);

        List<AbstractMessage<?>> messageSend = addMessage(ChatRequest.MessageType.USER, msg, image);
        request.setMessages(messageSend);

        YuanQiChatCompletion response = httpRequest(request);

        tokenTrack(response);
        log.info("Got response from {}.", getName());
        return response;
    }

    @Override
    YuanQiChatCompletion httpRequest(ChatRequest request) {
        Map<String, String> headers = new HashMap<>();
        headers.put("Authorization", "Bearer " + getAiAgentConfig().getToken());
        headers.put("date", ZonedDateTime.now().toString());
        headers.put("x-bce-date", ZonedDateTime.now().toString());
        return getRestTemplate().post(
                getAiAgentConfig().getUrl(),
                headers, request,
                YuanQiChatCompletion.class
        );
    }
}
